# -*- coding: utf-8 -*-
"""
@File: eval_gaitset.py
@Time: 2021/12/19 11:16
@Author: 鹄望潇湘
@desc:

Vgg16 使用随机的30帧，所有ID的数据进行训练30 epoch，使用ID为75-123的数据进行测试（网络输出维度为128），rank-1结果如下:
Evaluation complete. Cost: 0:00:57.076029
===Rank-1 (Include identical-view cases)===
NM: 86.718,	BG: 80.425,	CL: 69.775
===Rank-1 (Exclude identical-view cases)===
NM: 85.603,	BG: 79.205,	CL: 68.785
===Rank-1 of each angle (Exclude identical-view cases)===
NM: [80.20 81.84 88.47 87.65 85.61 86.43 88.37 87.14 90.20 82.86 82.86]
BG: [74.29 75.26 78.97 80.73 82.55 80.31 80.31 81.43 82.96 76.70 77.76]
CL: [63.47 63.47 73.47 71.84 72.65 70.51 72.75 70.61 69.90 67.25 60.71]
Vgg16 使用随机的30帧，0-74的ID的数据进行训练70 epoch，使用ID为75-123的数据进行测试（网络输出维度为124），rank-1结果如下:
Evaluation complete. Cost: 0:01:03.370398
===Rank-1 (Include identical-view cases)===
NM: 52.935,	BG: 35.316,	CL: 20.417
===Rank-1 (Exclude identical-view cases)===
NM: 49.156,	BG: 32.911,	CL: 19.546
===Rank-1 of each angle (Exclude identical-view cases)===
NM: [32.25 45.61 50.92 52.96 53.57 52.65 54.49 57.76 56.84 48.37 35.30]
BG: [26.33 31.75 32.58 33.64 40.61 33.57 31.63 31.12 34.08 35.67 31.02]
CL: [14.39 23.06 21.84 19.69 18.16 16.94 18.37 20.00 26.33 19.18 17.04]

CoAtNet 使用随机的30帧，0-74的ID的数据进行训练30 epoch，使用ID为75-123的数据进
行测试（网络输出维度为1024），rank-1结果如下:
Evaluation complete. Cost: 0:01:26.667511
===Rank-1 (Include identical-view cases)===
NM: 52.969,	BG: 35.092,	CL: 15.222
===Rank-1 (Exclude identical-view cases)===
NM: 49.072,	BG: 32.431,	CL: 14.555
===Rank-1 of each angle (Exclude identical-view cases)===
NM: [32.86 47.04 57.04 51.02 49.90 46.73 52.55 56.84 58.57 48.98 38.27]
BG: [27.04 33.81 42.37 33.85 32.35 26.23 25.31 32.55 36.53 36.29 30.41]
CL: [11.73 16.02 19.49 14.08 13.67 14.39 13.57 15.41 16.43 14.18 11.12]

CoAtNet 使用随机的30帧，0-74的ID的数据进行训练70 epoch，使用ID为75-123的数据进
行测试（网络输出维度为1024），rank-1结果如下:
Evaluation complete. Cost: 0:01:27.106347
===Rank-1 (Include identical-view cases)===
NM: 50.304,	BG: 36.524,	CL: 18.022
===Rank-1 (Exclude identical-view cases)===
NM: 45.993,	BG: 33.586,	CL: 17.180
===Rank-1 of each angle (Exclude identical-view cases)===
NM: [39.18 44.69 50.20 46.84 49.39 45.41 49.39 51.33 48.47 44.69 36.33]
BG: [31.74 37.01 40.41 34.79 30.10 28.06 30.31 33.47 34.69 38.15 30.71]
CL: [16.43 21.53 19.29 17.14 14.69 15.51 18.27 15.41 18.67 17.96 14.08]

"""
from datetime import datetime
import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm


from dataset.GaitMChannelDataSet import GaitMChannelDataLoader, GaitDataSet
from evaluator import evaluation


def boolean_string(s):
    if s.upper() not in {'FALSE', 'TRUE'}:
        raise ValueError('Not a valid boolean string')
    return s.upper() == 'TRUE'

# Exclude identical-view cases
def de_diag(acc, each_angle=False):
    result = np.sum(acc - np.diag(np.diag(acc)), 1) / 10.0
    if not each_angle:
        result = np.mean(result)
    return result


model = torch.load("../checkpoint_69.pth")

test_id = [index for index in range(75, 124)]
data_loader = GaitMChannelDataLoader("E:\\DataSet\\CASIA-B-preprocessed", True)
datas = data_loader.get_partial_data_with_id(test_id)
dataset = GaitDataSet(datas)
dataloader = DataLoader(dataset, batch_size=1, num_workers=0, shuffle=False)
gpu = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

time = datetime.now()
feature_list = list()
view_list = list()
seq_type_list = list()
label_list = list()
for id, item in enumerate(tqdm(dataloader)):
    label = item[1]
    angle = item[2][0]
    types = item[3][0]
    output = model(item[0].float().to(gpu))

    feature_list.append(output.cpu().detach().numpy())
    view_list.append(angle)
    seq_type_list.append(types)
    label_list.append(label)
test = (np.concatenate(feature_list, 0), view_list, seq_type_list, label_list)

print('Evaluating...')
acc = evaluation(test, {"dataset": "CASIA-B"})
print('Evaluation complete. Cost:', datetime.now() - time)

# Print rank-1 accuracy of the best model
# e.g.
# ===Rank-1 (Include identical-view cases)===
# NM: 95.405,     BG: 88.284,     CL: 72.041
for i in range(1):
    print('===Rank-%d (Include identical-view cases)===' % (i + 1))
    print('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
        np.mean(acc[0, :, :, i]),
        np.mean(acc[1, :, :, i]),
        np.mean(acc[2, :, :, i])))

# Print rank-1 accuracy of the best model，excluding identical-view cases
# e.g.
# ===Rank-1 (Exclude identical-view cases)===
# NM: 94.964,     BG: 87.239,     CL: 70.355
for i in range(1):
    print('===Rank-%d (Exclude identical-view cases)===' % (i + 1))
    print('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
        de_diag(acc[0, :, :, i]),
        de_diag(acc[1, :, :, i]),
        de_diag(acc[2, :, :, i])))


np.set_printoptions(precision=2, floatmode='fixed')
for i in range(1):
    print('===Rank-%d of each angle (Exclude identical-view cases)===' % (i + 1))
    print('NM:', de_diag(acc[0, :, :, i], True))
    print('BG:', de_diag(acc[1, :, :, i], True))
    print('CL:', de_diag(acc[2, :, :, i], True))
